#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
实用工具模块

功能：
1. 数据预处理与增强
2. 模型持久化工具
3. 结果可视化
4. 系统日志记录
"""

import json
import logging
import matplotlib.pyplot as plt
import numpy as np
import torch
import os
from datetime import datetime

from typing import Union, Optional, Dict, Any
import pandas as pd

def preprocess_data(
    data: Union[np.ndarray, pd.DataFrame],
    normalize: bool = True,
    min_max_scale: bool = False,
    clip_outliers: Optional[float] = None
) -> np.ndarray:
    """
    增强的数据预处理函数
    
    参数:
        data: 输入数据 (numpy数组或DataFrame)
        normalize: 是否进行Z-score归一化 (bool)
        min_max_scale: 是否进行Min-Max缩放 (bool)
        clip_outliers: 离群值截断阈值 (float)
        
    返回:
        处理后的数据 (numpy数组)
    """
    if not isinstance(data, np.ndarray):
        data = data.values if hasattr(data, 'values') else np.array(data)
    
    if clip_outliers is not None:
        mean, std = np.mean(data), np.std(data)
        data = np.clip(data, mean - clip_outliers*std, mean + clip_outliers*std)
    
    if normalize and min_max_scale:
        raise ValueError("不能同时使用normalize和min_max_scale")
    elif normalize:
        data = (data - np.mean(data)) / np.std(data)
    elif min_max_scale:
        data = (data - np.min(data)) / (np.max(data) - np.min(data))
        
    return data

def save_model(
    model: torch.nn.Module,
    path: str,
    model_name: str,
    metadata: Optional[Dict[str, Any]] = None,
    compress: bool = False
) -> None:
    """
    增强的模型保存函数
    
    参数:
        model: 要保存的模型
        path: 保存目录
        model_name: 模型名称
        metadata: 额外元数据
        compress: 是否压缩模型文件
    """
    try:
        if not os.path.exists(path):
            os.makedirs(path)
        
        save_path = os.path.join(path, f"{model_name}.pt")
        save_dict = {
            'state_dict': model.state_dict(),
            'metadata': metadata or {}
        }
        
        if compress:
            import zipfile
            zip_path = os.path.join(path, f"{model_name}.zip")
            with zipfile.ZipFile(zip_path, 'w') as zipf:
                with zipf.open(f"{model_name}.pt", 'w') as f:
                    torch.save(save_dict, f)
            logging.info(f"压缩模型已保存到 {zip_path}")
        else:
            torch.save(save_dict, save_path)
            logging.info(f"模型已保存到 {save_path}")
    except Exception as e:
        logging.error(f"保存模型失败: {e}")
        raise

def calculate_metrics(
    model: torch.nn.Module, 
    data_loader: torch.utils.data.DataLoader,
    metrics: list = ['accuracy']
) -> Dict[str, float]:
    """
    计算模型评估指标
    
    参数:
        model: 要评估的模型
        data_loader: 数据加载器
        metrics: 要计算的指标列表
        
    返回:
        指标名称到值的字典
    """
    model.eval()
    results = {}
    total = 0
    correct = 0
    total_loss = 0.0
    criterion = torch.nn.CrossEntropyLoss()
    
    with torch.no_grad():
        for inputs, targets in data_loader:
            outputs = model(inputs)
            _, predicted = torch.max(outputs.data, 1)
            total += targets.size(0)
            correct += (predicted == targets).sum().item()
            total_loss += criterion(outputs, targets).item()
    
    if 'accuracy' in metrics:
        results['accuracy'] = correct / total
    if 'loss' in metrics:
        results['loss'] = total_loss / len(data_loader)
        
    return results

def load_model(
    model: torch.nn.Module,
    path: str,
    model_name: str
) -> Dict[str, Any]:
    """
    增强的模型加载函数
    
    参数:
        model: 要加载参数的模型
        path: 模型目录
        model_name: 模型名称
        
    返回:
        包含模型状态和元数据的字典
    """
    try:
        load_path = os.path.join(path, f"{model_name}.pt")
        checkpoint = torch.load(load_path)
        model.load_state_dict(checkpoint['state_dict'])
        model.eval()
        logging.info(f"已从 {load_path} 加载模型")
        return checkpoint.get('metadata', {})
    except Exception as e:
        logging.error(f"加载模型失败: {e}")
        raise

def plot_results(x, y, title="结果可视化", xlabel="X轴", ylabel="Y轴"):
    """
    结果可视化函数
    
    参数:
        x: x轴数据
        y: y轴数据
        title: 图表标题
        xlabel: x轴标签
        ylabel: y轴标签
    """
    plt.figure(figsize=(10, 6))
    plt.plot(x, y, 'b-', linewidth=2)
    plt.title(title)
    plt.xlabel(xlabel)
    plt.ylabel(ylabel)
    plt.grid(True)
    plt.show()

def save_config(config, path="configs"):
    """
    保存配置到JSON文件
    
    参数:
        config: 配置字典
        path: 保存目录
    """
    if not os.path.exists(path):
        os.makedirs(path)
    
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    config_file = os.path.join(path, f"config_{timestamp}.json")
    
    with open(config_file, 'w') as f:
        json.dump(config, f, indent=4)
    
    print(f"配置已保存到 {config_file}")